'''
Created on 2010-8-20

@author: Madlee
'''

import numpy
import ctypes
from kuai.xyz import XYZ

kuaikan = numpy.ctypeslib.load_library(r'kuaikan', r'..\..\..\bin32w')

CTYPE_BYTE_TYPE = ctypes.c_byte

NUMPY_INDEX_TYPE = numpy.uint16
CTYPE_INDEX_TYPE = ctypes.c_ushort

NUMPY_FLOAT_TYPE = numpy.float32
CTYPE_FLOAT_TYPE = ctypes.c_float

NUMPY_BIG_INDEX_TYPE = numpy.uint32
CTYPE_BIG_INDEX_TYPE = ctypes.c_uint32

CTYPE_INDEX_POINTER = ctypes.POINTER(CTYPE_INDEX_TYPE)
CTYPE_FLOAT_POINTER = ctypes.POINTER(CTYPE_FLOAT_TYPE)
CTYPE_BYTES_POINTER = ctypes.POINTER(CTYPE_BYTE_TYPE)

def float_array(data):
    return numpy.array(data, dtype=NUMPY_FLOAT_TYPE)

def index_array(data):
    return numpy.array(data, dtype=NUMPY_INDEX_TYPE)

def byte_array(data):
    return numpy.array(data, dtype=numpy.uint8)


kuaikan.getConnectTableSize.restype = CTYPE_BIG_INDEX_TYPE
kuaikan.getConnectTableSize.argtypes = [CTYPE_INDEX_TYPE, CTYPE_INDEX_TYPE]

kuaikan.makeConnectTable.restype = None
kuaikan.makeConnectTable.argtypes = [CTYPE_INDEX_POINTER, CTYPE_INDEX_POINTER]

kuaikan.molRegion.restype = None
kuaikan.molRegion.argtypes = [CTYPE_INDEX_TYPE,     # nAtoms
                              CTYPE_FLOAT_POINTER,  # coords
                              CTYPE_FLOAT_POINTER,  # result
                              ]

kuaikan.drawMol3DLine.restype = None
kuaikan.drawMol3DLine.argtypes = [CTYPE_INDEX_POINTER,  # ct 
                                  CTYPE_FLOAT_POINTER,  # coords
                                  CTYPE_BYTES_POINTER   # colors
                                  ]

kuaikan.drawMol3DLine2.restype = None
kuaikan.drawMol3DLine2.argtypes = [CTYPE_INDEX_POINTER,  # ct 
                                  CTYPE_FLOAT_POINTER,  # coords
                                  CTYPE_BYTES_POINTER   # colors
                                  ]

kuaikan.setAxis.restype = None
kuaikan.setAxis.argtypes = [CTYPE_FLOAT_POINTER]

kuaikan.rotateAxis.restype = None
kuaikan.rotateAxis.argtypes = [CTYPE_FLOAT_POINTER, CTYPE_FLOAT_TYPE, CTYPE_FLOAT_TYPE]

kuaikan.kuaikanAtomBalls.restype = None
kuaikan.kuaikanAtomBalls.argtypes = [CTYPE_INDEX_TYPE, CTYPE_FLOAT_POINTER, CTYPE_BYTES_POINTER, CTYPE_FLOAT_POINTER, CTYPE_INDEX_POINTER, CTYPE_INDEX_TYPE]

def getConnectTableSize(natoms, nbonds):
    return kuaikan.getConnectTableSize(natoms, nbonds)

def makeConnectTable(data):
    nsize = getConnectTableSize(data[0], data[1])
    result = numpy.zeros((nsize,), dtype=NUMPY_INDEX_TYPE)
    input = numpy.array(data, dtype=NUMPY_INDEX_TYPE)
    kuaikan.makeConnectTable(input.ctypes.data_as(CTYPE_INDEX_POINTER), result.ctypes.data_as(CTYPE_INDEX_POINTER))
    return result


def drawMol3DLine(ct, coords, colors):
    kuaikan.drawMol3DLine(ct.ctypes.data_as(CTYPE_INDEX_POINTER),
                          coords.ctypes.data_as(CTYPE_FLOAT_POINTER),
                          colors.ctypes.data_as(CTYPE_BYTES_POINTER))

def drawMol3DLine2(ct, coords, colors):
    kuaikan.drawMol3DLine2(ct.ctypes.data_as(CTYPE_INDEX_POINTER),
                          coords.ctypes.data_as(CTYPE_FLOAT_POINTER),
                          colors.ctypes.data_as(CTYPE_BYTES_POINTER))

def drawAtomBalls(natoms, coords, colors, radius, atoms=None):
    if atoms:
        kuaikan.kuaikanAtomBalls(natoms, coords.ctypes.data_as(CTYPE_FLOAT_POINTER),
                              colors.ctypes.data_as(CTYPE_BYTES_POINTER), 
                              radius.ctypes.data_as(CTYPE_FLOAT_POINTER),
                              atoms.ctypes.data_as(CTYPE_INDEX_POINTER), 6)
    else:
        kuaikan.kuaikanAtomBalls(natoms, coords.ctypes.data_as(CTYPE_FLOAT_POINTER),
                              colors.ctypes.data_as(CTYPE_BYTES_POINTER), 
                              radius.ctypes.data_as(CTYPE_FLOAT_POINTER),
                              atoms, 6)
                              
        
        
def mol_region(coords):
    result = float_array([0] * 6)
    kuaikan.molRegion(len(coords)/3, 
                      coords.ctypes.data_as(CTYPE_FLOAT_POINTER),
                      result.ctypes.data_as(CTYPE_FLOAT_POINTER))
    p1 = XYZ(float(result[0]), float(result[1]), float(result[2])) 
    p2 = XYZ(float(result[3]), float(result[4]), float(result[5]))
    size = p2-p1
    radius = abs(size)/2
    return p1, p2, (p1+p2)/2, size, radius
    
def setAxis(arg):
    kuaikan.setAxis(arg.ctypes.data_as(CTYPE_FLOAT_POINTER))

def rotateAxis(arg, x, y):
    kuaikan.rotateAxis(arg.ctypes.data_as(CTYPE_FLOAT_POINTER), x, y)
    
